import numpy as np
from view.endSrc.WeightedObject import WeightedObject

from view.endSrc.DBConfig import DBConfig


class ProbClusters(np.ndarray, WeightedObject):

    def __new__(cls, probMatrix=None, labels=None, weight=float('inf')):

        assert (probMatrix is not None) or \
               (labels is not None), 'probMatrix and labels can not be both None: ' \
                + str(probMatrix) + ', ' + str(labels)

        if probMatrix is not None:
            # check shape is 2D
            assert len(probMatrix.shape) == 2, 'probMatrix must be like (n, k)'

            # Input array is an already formed ndarray instance,
            # we first cast to be our class type
            self = np.asarray(probMatrix).view(cls)

            # set properties
            self.nRows = probMatrix.shape[0]
            self.nCols = probMatrix.shape[1]
            labels = probMatrix.argmax(axis=1)

        elif labels is not None:

            assert isinstance(labels, list) or \
                   (isinstance(labels, np.ndarray) and len(labels.shape) == 1), \
                'labels must be list or np.ndarray with shape like (d, )'

            # to integer ids and check validation
            labels = np.asarray(labels).astype(int)
            nClusters = len(np.unique(labels))
            for b in labels:
                assert b < nClusters and b >= 0, 'all labels must be in [0, ' + str(nClusters) + ')'

            # new obj
            self = np.asarray(float('inf')).view(cls)

            # set properties
            self.nRows = len(labels)
            self.nCols = nClusters


        # add labels, clusterIdsList property
        self.labels = labels
        return self

    def __init__(self, probMatrix=None, labels=None, weight=float('inf')):
        '''
            probMatrix, labels must be same as that of __new__,
            but we do not use it, since they had been already used in
            __new__ to initialize the properties like self.nRows, self.clusterIdsList, etc.
        :param probMatrix:
        :param labels:
        '''

        # invoke the superclass WeightedObject's __init__
        # https://stackoverflow.com/questions/9575409/calling-parent-class-init-with-multiple-inheritance-whats-the-right-way
        # set logLikelihood to the smallest one, so that
        # it will not be chosen by maximizing logLikelihood method
        WeightedObject.__init__(self, weight)

        self.clusterIdsList = []
        for i in range(self.nCols):
            self.clusterIdsList.append([])
        for i, cId in enumerate(self.labels):
            self.clusterIdsList[cId].append(i)

    def __str__(self):
        '''
            for simple test
        :return:
        '''
        tmp = 'weight: ' + str(self.weight) + '\n'
        tmp += 'shape: ' + str(self.shape) + '\n'
        tmp += 'nClusters: ' + str(self.getNumOfClusters()) + '\n'
        tmp += 'nRows: ' + str(self.nRows) + '\n'
        tmp += 'nCols: ' + str(self.nCols) + '\n'
        tmp += 'labels: ' + str(self.labels) + '\n'
        tmp += 'clusterIdsList: ' + str(self.clusterIdsList) + '\n'
        tmp += 'probMat: ' + str(np.ndarray.__str__(self)) + '\n'
        tmp += 'labels: ' + str(self.labels) + '\n'
        return tmp

    def getIds(self, clusterId):
        '''

        :param clusterId: in [0, self.nClos)
        :return: list<int> for the ids of the cluster whose id = clusterId
        '''
        assert clusterId >= 0 and clusterId < self.nCols, 'clusterId=' + \
            str(clusterId) + ' must be positive and less than ' + str(self.nCols)

        return self.clusterIdsList[clusterId]

    def getIdsOfEachCluster(self):
        '''

            clusterIdsList: a list containing k lists of indexes that belong to the cluster.
                For example, if [0, 0, 1, 2, 0, 2, 1] is labels,
                    then clusterIdsList is [[0, 1, 4], [2, 6], [3, 5]]
        :return:  clusterIdsList
        '''

        return self.clusterIdsList

    def getNumOfClusters(self):
        return self.nCols

    def getLabels(self):
        '''
            return the list<int>, i.e., the labels that are clusterIds of all records
        :return:
        '''
        return self.labels

    def __getitem__(self, item):
        '''

        :param item: (i, j), where i is index of record, j is the clusterId
        :return: the probility of the i-th record belonging to the j-th cluster
        '''

        if isinstance(item, int) or \
           (isinstance(item, tuple) and len(item) != 2) or \
           len(self.shape) != 0:
            return super().__getitem__(item)

        # assert len(item) == 2, 'probility index must be like (i,j)'

        i, j = item[0], item[1]

        # if len(self.shape) == 2:
        #     return super().__getitem__((i,j))

        return 1.0 if self.labels[i] == j else 0.0


    def saveToFile(self, pathFileName):
        '''
            Save to self.labels or self(probMatrix) to files
        :param pathFileName:
        :return: True if ok, False otherwise.
        '''

        logger = DBConfig().logger

        # is probMatrix available ?
        if len(self.shape) == 2:
            try:
                # write np.ndarray_1D to disk, so that we can load back with
                # dimension kept.
                np.savetxt(pathFileName, self, delimiter=",")
            except Exception as e:
                logger.write('ProbClusters.saveToFile..failed to write probMatrix to file: ' + pathFileName)
                return False
            return True

        # otherwise, save labels to file to save the disk space.
        try:
            # write np.ndarray_1D to disk, so that we can load back with
            # dimension kept.
            np.savetxt(pathFileName, np.asarray(self.labels), delimiter=",")
        except Exception as e:
            logger.write('ProbClusters.saveToFile..failed to write labels to file: ' + pathFileName)
        return True


    def getDataDimension(self):
        '''
            The data dimension denotes the dimension of records of a dataset.
            The data dimension can be derived from the parameter dimension of the subclasses of ProbClusters.

            For example, for GaussianClusters, the data dimension is the dimension of the mean vector parameter
            for GammaClusters the data dimension is the dimension of the shape vector parameter.
        :return: the dimension of records of a dataset,
                whose cluster information is collected by the subclasses of ProbClusters
        '''

        # https://docs.python.org/3/library/exceptions.html#exception-hierarchy
        assert False, 'getDataDimension must be implemented by subclasses of ProbClusters'


    def dumpParamsToJSONDict(self):
        '''
            This method is used to get the JSON dict representation of all parameters,
            used to generate a synthetic dataset, such as Guassian mixtures, Gamma mixtures,
            so that user can write the JSON dict representation to DB

        :return: JSON dict represenation for all parameters, where
                it must pass the MyTools.isJsonDumpable testing.
        '''

        assert False, 'dumpParamsToJSONDict must be implemented by subclasses of ProbClusters'


    @staticmethod
    def testIsInstance():
        x = np.asarray([[0.2, 0.3, 0.5],
                        [0.1, 0.6, 0.3],
                        [0.5, 0.3, 0.1],
                        [0.6, 0.2, 0.2]])

        pc = ProbClusters(x)
        print('pc.weight', pc.weight)

        print('is WeightedObject instance: ', isinstance(pc, WeightedObject))
        print('is np.ndarray instance: ', isinstance(pc, np.ndarray))
        print('pc.shape: ', pc.shape)

        ids = pc.getIdsOfEachCluster()  # [[2, 3], [1], [0]]

        print(ids)
        print('pc.weight', pc.weight)

        pc.weight = 10.0
        print('pc.weight: ', pc.weight)

        print('pc[0,1]: ', pc[0, 1])

    @staticmethod
    def testLabels():
        x = np.asarray([[0.2, 0.3, 0.5],
                        [0.1, 0.6, 0.3],
                        [0.5, 0.3, 0.1],
                        [0.6, 0.2, 0.2]])

        x = np.argmax(x, axis=1)

        print('labels: ', x)

        pc = ProbClusters(labels=x)
        print('pc.weight', pc.weight)

        print('is WeightedObject instance: ', isinstance(pc, WeightedObject))
        print('is np.ndarray instance: ', isinstance(pc, np.ndarray))
        print('pc.shape: ', pc.shape)

        ids = pc.getIdsOfEachCluster()  # [[2, 3], [1], [0]]

        print(ids)
        print('pc.weight', pc.weight)

        pc.weight = 10.0
        print('pc.weight: ', pc.weight)

        print('pc.nClusters: ', pc.getNumOfClusters())


    @staticmethod
    def testGetItem():
        x = np.asarray([[0.2, 0.3, 0.5],
                        [0.1, 0.6, 0.3],
                        [0.5, 0.3, 0.1],
                        [0.6, 0.2, 0.2]])

        y = np.argmax(x, axis=1)
        print('labels: ', y)

        pc = ProbClusters(probMatrix=x)
        print(pc[0,1])

        pcy = ProbClusters(labels=y)
        print(pcy[0, 1])  # 0
        print(pcy[0, 2])  # 1.0


    @staticmethod
    def testPrint():
        x = np.asarray([[0.2, 0.3, 0.5],
                        [0.1, 0.6, 0.3],
                        [0.5, 0.3, 0.1],
                        [0.6, 0.2, 0.2]])

        lab = np.argmax(x, axis=1)

        # print('labels: ', lab)  # [2 1 0 0]

        pc = ProbClusters(labels=lab)
        print(pc)

        print('-------------------------')
        pmat = ProbClusters(probMatrix=x)
        print(pmat)


if __name__ == '__main__':

    # ProbClusters.testIsInstance()  # ok
    # ProbClusters.testLabels()
    ProbClusters.testGetItem()
    # ProbClusters.testPrint()
